perf: optimize fused AllReduce + RMSNorm (custom_all_reduce) - #3817

Motivation

This PR makes three small updates to the custom fused AllReduce + RMSNorm kernels:

• use the existing DPP-based block reduce on the one-stage decode path;
• simplify the two-stage stage-1 reduce-scatter path to avoid LDS and block-level sync;
• move the two-stage stage-2 RMSNorm reduction to the shared DPP wave-reduce primitive.

Technical Details

1stage

• decode path. ar_fusion_epilogue_block_reduce now uses multithread_reduce from hip_reduce.h instead of the generic warpReduce implementation. This switches the intra-wave reduction from the ds_bpermute XOR butterfly to the existing DPP tree reduce.

2stage

• reduce-scatter. reduce_scatter_cross_device_store is rewritten as a flat reduce-scatter. Each thread reduces one pack across all input ranks in fp32, downcasts the result to bf16, and writes it to every rank's temporary buffer directly. This removes the previous LDS staging and __syncthreads() from the warp-per-rank implementation. The reduce order over ptrs[0..ngpus-1] is preserved, so the result remains bit-identical to the previous path.

• stage-1 launch config. The stage-1 launch now uses block size 256 instead of 512. This gives small-m decode shapes more parallelism across CUs.

• stage-2 RMSNorm reduce. local_device_load_rmsnorm now uses wave_reduce from hip_reduce.h for the RMSNorm square-sum reduction. This keeps the reduction implementation consistent with the shared DPP primitive. In testing, this change alone is performance-neutral.

This PR does not change the one-stage/two-stage selection heuristic. The thresholds discussed in #3458 can therefore still be used as-is.

Test Plan

Test environment:

• GPU: MI355X
• ROCm: 7.2.3

Benchmark:

The benchmark script (test_fused_ar_rmsnorm_perf.py is inspired by the script posted by @TennyWang1223 in #3458, with a few changes to improve stability.

To reduce host launch overhead and cross-rank rendezvous jitter, 2000 op calls are captured into one CUDA graph, and a single replay is timed with CUDA events. The reported latency is the minimum over 80 replays, taking the max latency across ranks. GPU clocks were pinned with:

rocm-smi --setperfdeterminism 2400

Reproduce commands:

TP_LIST={4,8} AITER_AR_1STAGE={0,1} PYTHONPATH=$PWD \
  python <script_path>

Restore GPU clocks after benchmarking:

for g in 0 1 2 3 4 5 6 7; do
  rocm-smi --setperflevel auto -d $g
done

Test Results

For the CDF plots below, curves closer to the upper-left corner indicate better performance.

One-stage path

• TP4

• TP8

Two-stage path

• TP4

• TP8

Under the selection thresholds from #3458 (one-stage for m <= 32 at TP4 and m <= 16 at TP8), the latency-critical small-m decode shapes are served by the one-stage path, which this PR speeds up. The shapes where the two-stage path regresses largely fall in that one-stage range, so the practical impact is minimal; on the larger shapes that do use the two-stage path the change is a 2.5% net gain.

This PR does not change the 1stage/2stage selection heuristic, so the threshold from #3458 remains applicable.